package crypt

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"os"

	"gitee.com/851537/gbase"
)

type Rsa struct {
	Base
	PublicKey  *rsa.PublicKey
	PrivateKey *rsa.PrivateKey
}

// 生成RSA密钥对
func (this *Rsa) GenerateKeyPair(bitSize int) (err error) {
	// 生成RSA密钥对
	this.PrivateKey, err = rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return fmt.Errorf("生成RSA密钥对失败: %w", err)
	}
	this.PublicKey = &this.PrivateKey.PublicKey
	return nil
}

// 保存私钥到文件
func (this *Rsa) SavePrivateKeyToFile(filePath string) error {
	// 检查私钥是否存在
	if !this.HasPrivateKey() {
		return errors.New("私钥不存在")
	}
	// 编码私钥
	privKeyBytes := x509.MarshalPKCS1PrivateKey(this.PrivateKey)
	// 创建PEM块
	privKeyPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privKeyBytes,
	}
	// 写入文件
	err := os.WriteFile(filePath, pem.EncodeToMemory(privKeyPEM), 0600)
	if err != nil {
		return fmt.Errorf("写入私钥文件失败: %w", err)
	}
	if this.privateKeyFile == nil {
		this.privateKeyFile = gbase.NewFile()
	}
	this.privateKeyFile.Path = filePath
	return nil
}

// 保存公钥到文件
func (this *Rsa) SavePublicKeyToFile(filePath string) error {
	// 检查公钥是否存在
	if !this.HasPublicKey() {
		return errors.New("公钥不存在")
	}
	// 编码公钥
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(this.PublicKey)
	if err != nil {
		return fmt.Errorf("编码公钥失败: %w", err)
	}
	// 创建PEM块
	pubKeyPEM := &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubKeyBytes,
	}
	// 写入文件
	err = os.WriteFile(filePath, pem.EncodeToMemory(pubKeyPEM), 0644)
	if err != nil {
		return fmt.Errorf("写入公钥文件失败: %w", err)
	}
	if this.publicKeyFile == nil {
		this.publicKeyFile = gbase.NewFile()
	}
	this.publicKeyFile.Path = filePath
	return nil
}

func (this *Rsa) LoadPrivateKey(privateKeyData []byte) error {
	// 解析PEM格式
	block, _ := pem.Decode(privateKeyData)
	if block == nil {
		return errors.New("私钥文件格式错误：非PEM格式")
	}

	// 检查PEM块类型
	if block.Type != "PRIVATE KEY" && block.Type != "RSA PRIVATE KEY" {
		return fmt.Errorf("私钥类型错误: %s", block.Type)
	}

	// 解析私钥
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		// 尝试解析PKCS8格式
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return fmt.Errorf("解析私钥失败: %w", err)
		}
		rsaPrivateKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return errors.New("私钥不是RSA类型")
		}
		privateKey = rsaPrivateKey
	}

	this.PrivateKey = privateKey
	return nil
}

func (this *Rsa) LoadPrivateKeyFromFile(filePath string) error {
	// 读取私钥文件
	privKeyData, err := os.ReadFile(filePath)
	if err != nil {
		return fmt.Errorf("读取私钥文件失败: %w", err)
	}
	return this.LoadPrivateKey(privKeyData)
}

func (this *Rsa) LoadPublicKey(publicKeyData []byte) error {
	// 解析PEM格式
	block, _ := pem.Decode(publicKeyData)
	if block == nil {
		return errors.New("公钥文件格式错误：非PEM格式")
	}

	// 检查PEM块类型
	if block.Type != "PUBLIC KEY" && block.Type != "RSA PUBLIC KEY" {
		return fmt.Errorf("公钥类型错误: %s", block.Type)
	}

	// 解析公钥
	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		// 尝试解析PKCS1格式
		publicKey, err = x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return fmt.Errorf("解析公钥失败: %w", err)
		}
	}

	// 转换为rsa.PublicKey类型
	rsaKey, ok := publicKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("公钥不是RSA类型")
	}
	this.PublicKey = rsaKey
	return nil
}

func (this *Rsa) LoadPublicKeyFromFile(filePath string) error {
	// 读取公钥文件
	pubKeyData, err := os.ReadFile(filePath)
	if err != nil {
		return fmt.Errorf("读取公钥文件失败: %w", err)
	}
	return this.LoadPublicKey(pubKeyData)
}

func (this *Rsa) HasPrivateKey() bool {
	return this.PrivateKey != nil
}

func (this *Rsa) HasPublicKey() bool {
	return this.PublicKey != nil
}

func (this *Rsa) Sign(payload []byte) (sign []byte, err error) {
	if !this.HasPrivateKey() && this.privateKeyFile.Path == "" {
		return nil, errors.New("未设置私钥或私钥文件")
	}
	if !this.HasPrivateKey() {
		if err = this.LoadPrivateKeyFromFile(this.privateKeyFile.Path); err != nil {
			return nil, err
		}
	}

	hash := sha1.New()
	hash.Write(payload)
	hashSum := hash.Sum(nil)

	sign, err = rsa.SignPKCS1v15(rand.Reader, this.PrivateKey, crypto.SHA1, hashSum)
	if err != nil {
		return nil, fmt.Errorf("签名失败: %w", err)
	}
	return sign, nil
}

func (this *Rsa) SignString(payload string) (sign string, err error) {
	var bsign []byte
	bsign, err = this.Sign([]byte(payload))
	if err != nil {
		return "", err
	}
	sign = string(bsign)
	return
}

func (this *Rsa) SignB64(payload string) (sign string, err error) {
	var tmpSign []byte
	tmpSign, err = this.Sign([]byte(payload))
	if err != nil {
		return "", err
	}
	sign = base64.StdEncoding.EncodeToString(tmpSign)
	return
}

func (this *Rsa) Verify(sign []byte, payload []byte) (err error) {
	if !this.HasPublicKey() && this.publicKeyFile.Path == "" {
		return errors.New("未设置公钥或公钥文件")
	}
	if !this.HasPublicKey() {
		if err = this.LoadPublicKeyFromFile(this.publicKeyFile.Path); err != nil {
			return err
		}
	}

	hash := sha1.New()
	hash.Write(payload)
	hashSum := hash.Sum(nil)

	err = rsa.VerifyPKCS1v15(this.PublicKey, crypto.SHA1, hashSum, sign)
	if err != nil {
		return fmt.Errorf("签名验证失败: %w", err)
	}
	return nil
}

func (this *Rsa) VerifyB64(sign string, payload string) (err error) {
	signBytes, err := base64.StdEncoding.DecodeString(sign)
	if err != nil {
		return fmt.Errorf("解码签名失败: %w", err)
	}
	return this.Verify(signBytes, []byte(payload))
}

// 使用RSA公钥加密数据
func (this *Rsa) Encrypt(content []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, this.PublicKey, content)
}

// 使用RSA私钥加密数据
func (this *Rsa) Decrypt(content []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, this.PrivateKey, content)
}

func NewRsa(opts ...gbase.Option[*Rsa]) *Rsa {
	obj := new(Rsa)
	for _, v := range opts {
		v(obj)
	}
	return obj
}

// 直接提供公私钥字符串创建Rsa对象
func NewRsaWithPem(privateKey, publicKey string) (obj *Rsa, err error) {
	obj = NewRsa()
	if privateKey != "" {
		if err = obj.LoadPrivateKey([]byte(privateKey)); err != nil {
			return
		}
	}
	if publicKey != "" {
		if err = obj.LoadPublicKey([]byte(publicKey)); err != nil {
			return
		}
	}
	return obj, nil
}
